#include <utility>
#include <stdio.h>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <atomic>
#include "net/tcp_client.hpp"
#include "net/io_service_pool.hpp"
#include "base/logger.h"

using namespace easyasio::net;
class Client;

class Session 
{
public:
    Session(asio::io_service& loop, const std::string& remote_addr, int remote_port, Client* owner)
        : client_(std::make_shared<TcpClient>(loop, remote_addr, remote_port)), owner_(owner), bytesRead_(0), bytesWritten_(0), messagesRead_(0)
    {
        client_->setConnectionCallback(
            std::bind(&Session::onConnection, this, std::placeholders::_1));
        client_->setMessageCallback(
            std::bind(&Session::onMessage, this, std::placeholders::_1, std::placeholders::_2));
    }

    void start() { client_->asynConnect(); }
    void stop() { client_->stop(); }
    int64_t bytesRead() const { return bytesRead_; }
    int64_t messagesRead() const { return messagesRead_; }

private:
    void onConnection(const TcpConnectionPtr& conn);
    void onDisconnect(const TcpConnectionPtr& conn);
    void onMessage(const TcpConnectionPtr& conn, Buffer* buf)
    {
        if (!conn || !conn->connected())
            return;

        ++messagesRead_;
        bytesRead_ += buf->readableBytes();
        bytesWritten_ += buf->readableBytes();
        std::string msg(buf->peek(), buf->readableBytes());
        conn->send(msg);
        buf->retrieveAll();
    }

private:
    std::shared_ptr<TcpClient> client_;
    Client* owner_;
    int64_t bytesRead_;
    int64_t bytesWritten_;
    int64_t messagesRead_;
};

class Client
{
 public:
  Client(const std::string& remote_addr, int remote_port, int blockSize,
         int sessionCount, int timeout, int threadCount)
    : threadPool_(threadCount),
      loop_timer_(threadPool_.get_io_service()),
      sessionCount_(sessionCount),
      timeout_(timeout)
{
    numConnected_ = 0;
    loop_timer_.expires_from_now(std::chrono::seconds(timeout));
    loop_timer_.async_wait(std::bind(&Client::handleTimeout, this));

    for (int i = 0; i < blockSize; ++i)
    {
        message_.push_back(static_cast<char>(i % 128));
    }

    for (int i = 0; i < sessionCount; ++i)
    {
        char buf[32];
        #ifdef WIN32
        sprintf_s(buf, sizeof buf, "C%05d", i);
        #else
        snprintf(buf, sizeof buf, "C%05d", i);
        #endif

        std::shared_ptr<Session> session = std::make_shared<Session>(threadPool_.get_io_service(), remote_addr, remote_port, this);
        session->start();
        sessions_.push_back(session);
    }
}

const std::string& message() const
{
    return message_;
}

void onConnect()
{
    ++numConnected_;
    if (numConnected_ == sessionCount_)
    {
        LOG_WARN("all connected");
    }
}

void run()
{
    threadPool_.run();
}

void onDisconnect()
{
    --numConnected_;
    if (numConnected_ == 0)
    {
        LOG_WARN("all disconnected");

        int64_t totalBytesRead = 0;
        int64_t totalMessagesRead = 0;
        for (auto it = sessions_.begin(); it != sessions_.end(); ++it)
        {
            totalBytesRead += (*it)->bytesRead();
            totalMessagesRead += (*it)->messagesRead();
        }
        LOG_WARN("{} total bytes read", totalBytesRead);
        LOG_WARN("{} total messages read", totalMessagesRead);
        LOG_WARN("{} average message size", static_cast<double>(totalBytesRead) / static_cast<double>(totalMessagesRead));
        LOG_WARN("{} MiB/s throughput", static_cast<double>(totalBytesRead) / (timeout_ * 1024 * 1024));
        threadPool_.stop();
    }
}

private:
    void handleTimeout()
    {
        LOG_WARN("stop");
        std::for_each(sessions_.begin(), sessions_.end(), std::mem_fn(&Session::stop));
    }

private:
    io_service_pool threadPool_;
    asio::basic_waitable_timer<std::chrono::system_clock> loop_timer_;
    int sessionCount_;
    int timeout_;

    std::vector<std::shared_ptr<Session>> sessions_;
    std::string message_;
    std::atomic_int numConnected_;
    std::shared_ptr<std::thread> iopool_thread_;
};

void Session::onDisconnect(const TcpConnectionPtr& conn)
{
    owner_->onDisconnect();
}

void Session::onConnection(const TcpConnectionPtr& conn)
{
    if (conn->connected())
    {
        conn->setTcpNoDelay(true);
        conn->send(owner_->message());
        owner_->onConnect();
        conn->setCloseCallback(std::bind(&Session::onDisconnect, this, std::placeholders::_1));
    }
    else
    {
        owner_->onDisconnect();
    }
}

int main(int argc, char* argv[])
{
    if (argc != 7)
    {
        fprintf(stderr, "Usage: client <host_ip> <port> <threads> <blocksize> ");
        fprintf(stderr, "<sessions> <time>\n");
    }
    else
    {
        easyasio::base::Logger::instance().setLevel(easyasio::base::Logger::LOGLEVEL_WARN);
        const char* ip = argv[1];
        uint16_t port = static_cast<uint16_t>(atoi(argv[2]));
        int threadCount = atoi(argv[3]);
        int blockSize = atoi(argv[4]);
        int sessionCount = atoi(argv[5]);
        int timeout = atoi(argv[6]);

        Client client(ip, port, blockSize, sessionCount, timeout, threadCount);
        client.run();
    }

    return 0;
}

